import json
import sys
import datetime
import hashlib
import hmac
import requests
import time

method = "POST"
host = "visual.volcengineapi.com"
region = "cn-north-1"
endpoint = "https://visual.volcengineapi.com"
service = "cv"

def sign(key, msg):
    return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()

def getSignatureKey(key, dateStamp, regionName, serviceName):
    kDate = sign(key.encode("utf-8"), dateStamp)
    kRegion = sign(kDate, regionName)
    kService = sign(kRegion, serviceName)
    kSigning = sign(kService, "request")
    return kSigning

def formatQuery(parameters):
    request_parameters_init = ""
    for key in sorted(parameters):
        request_parameters_init += key + "=" + parameters[key] + "&"
    request_parameters = request_parameters_init[:-1]
    return request_parameters

def signV4Request(access_key, secret_key, service, req_query, req_body):
    if access_key is None or secret_key is None:
        print("No access key is available.")
        sys.exit()

    t = datetime.datetime.utcnow()
    current_date = t.strftime("%Y%m%dT%H%M%SZ")

    datestamp = t.strftime("%Y%m%d")  # Date w/o time, used in credential scope
    canonical_uri = "/"
    canonical_querystring = req_query
    signed_headers = "content-type;host;x-content-sha256;x-date"
    payload_hash = hashlib.sha256(req_body.encode("utf-8")).hexdigest()
    content_type = "application/json"
    canonical_headers = (
        "content-type:"
        + content_type
        + "\n"
        + "host:"
        + host
        + "\n"
        + "x-content-sha256:"
        + payload_hash
        + "\n"
        + "x-date:"
        + current_date
        + "\n"
    )
    canonical_request = (
        method
        + "\n"
        + canonical_uri
        + "\n"
        + canonical_querystring
        + "\n"
        + canonical_headers
        + "\n"
        + signed_headers
        + "\n"
        + payload_hash
    )

    algorithm = "HMAC-SHA256"
    credential_scope = datestamp + "/" + region + "/" + service + "/" + "request"
    string_to_sign = (
        algorithm
        + "\n"
        + current_date
        + "\n"
        + credential_scope
        + "\n"
        + hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()
    )
    signing_key = getSignatureKey(secret_key, datestamp, region, service)
    signature = hmac.new(
        signing_key, (string_to_sign).encode("utf-8"), hashlib.sha256
    ).hexdigest()

    authorization_header = (
        algorithm
        + " "
        + "Credential="
        + access_key
        + "/"
        + credential_scope
        + ", "
        + "SignedHeaders="
        + signed_headers
        + ", "
        + "Signature="
        + signature
    )

    headers = {
        "X-Date": current_date,
        "Authorization": authorization_header,
        "X-Content-Sha256": payload_hash,
        "Content-Type": content_type,
    }

    # ************* SEND THE REQUEST *************
    request_url = endpoint + "?" + canonical_querystring

    try:
        print(request_url, '\n', headers, '\n', req_body)
        r = requests.post(request_url, headers=headers, data=req_body)
    except Exception as err:
        print(f"error occurred: {err}")
        raise
    else:
        resp_str = r.text.replace("\\u0026", "&")
        dic = json.loads(resp_str)
        assert dic is not None, dic
        return dic

def creat_task(input_image_urls, prompt, height=1024, width=1024):
    # 请求Query，按照接口文档中填入即可
    query_params = {
        "Action": "CVSync2AsyncSubmitTask",
        "Version": "2022-08-31",
    }
    formatted_query = formatQuery(query_params)

    body_params = {
        "req_key": "jimeng_t2i_v40",
        "image_urls": input_image_urls,
        "prompt": prompt,
        "force_single": True,
    }
    formatted_body = json.dumps(body_params)

    dic = signV4Request(
        access_key, secret_key, service, formatted_query, formatted_body
    )
    print('==========dic==============')
    print(dic)
    return dic

def get_result(task_id):
    query_params = {
        "Action": "CVSync2AsyncGetResult",
        "Version": "2022-08-31",
    }
    formatted_query = formatQuery(query_params)

    req_json_dict = {"return_url": True}

    req_json_str = json.dumps(req_json_dict, ensure_ascii=False)

    body_params = {
        "req_key": "jimeng_t2i_v40",
        "task_id": task_id,
        "req_json": req_json_str,
    }

    formatted_body = json.dumps(body_params, ensure_ascii=False)

    dic = signV4Request(
        access_key, secret_key, service, formatted_query, formatted_body
    )
    return dic

def get_result_with_polling(task_id, max_attempts=60, interval=3):
    """
    轮询获取异步任务结果，直到状态为 'done' 或超时
    :param task_id: 任务ID
    :param max_attempts: 最大轮询次数
    :param interval: 每次轮询间隔（秒）
    :return: 最终响应字典（包含图片URL等）
    """
    for attempt in range(1, max_attempts + 1):
        print(f"🔁 第 {attempt}/{max_attempts} 次查询任务 {task_id} 状态...")

        try:
            dic = get_result(task_id)  # 调用你写的函数
        except Exception as e:
            print(f"⚠️ 查询异常: {e}")
            time.sleep(interval)
            continue

        # 获取任务状态
        status = dic.get("data", {}).get("status", "")
        print(f"📊 当前状态: {status}")

        if status == "done":
            print("✅ 任务已完成！")
            return dic  # 返回最终结果，包含 image_url 等
        elif status == "in_queue":
            print("⏳ 任务排队中，等待中...")
        else:
            print(f"⚠️ 未知状态: {status}")

        # 如果不是 done，等待后重试
        if attempt < max_attempts:
            time.sleep(interval)
        else:
            print("⏰ 轮询超时，未获取到完成状态")

    return None  # 超时返回 None 或抛异常，根据你的逻辑决定

if __name__ == "__main__":
    # 请求凭证，从访问控制申请
    import time

    access_key = 'AKLTY2UwODFmNmU1MzRjNDE5MTg4ZmM2ZmJiMzIzMDRiNjE'
    secret_key = 'WXpRMllXSXdNRE5rT1RRd05HWXdOR0k1T1RWak0yUm1ZV1psT1dNeVlqZw=='
    image_urls = r"https://huayiyi-ai.oss-cn-shanghai.aliyuncs.com/test_tmp%2F1756904274434%20-%20bfc04789-9331-4993-9f25-49ee0ad22ceb.jpg"
    prompt = "将这件长款衣服裁剪为短款，仅缩短下摆部分，保持上半身设计、版型和细节完全不变。"
    # prompt = "\u5c06\u8fd9\u4ef6\u957f\u6b3e\u8863\u670d\u88c1\u526a\u4e3a\u77ed\u6b3e\uff0c\u4ec5\u7f29\u77ed\u4e0b\u6446\u90e8\u5206\uff0c\u4fdd\u6301\u4e0a\u534a\u8eab\u8bbe\u8ba1\u3001\u7248\u578b\u548c\u7ec6\u8282\u5b8c\u5168\u4e0d\u53d8\u3002"
    


    task_ids = []
    succ = fail = 0
    for i in range(1,2):
        print('😁第',i,'次 获取结果')

        dic = creat_task([image_urls], prompt)
        task_id = dic["data"]["task_id"]
        # time.sleep(3)

        task_ids.append(task_id)

    print("开始获取结果")
    for task_id in task_ids:
        # task_id = 13750662383678537459
        dic = get_result_with_polling(task_id)
        # print(dic)

        if dic["data"]["image_urls"] is not None:
            print(dic['data']['image_urls'][0])
            succ += 1
        else:
            print("未获取到图片URL")
            fail += 1

        # time.sleep(0.5)

    print('succ ', succ)
    print('fail ', fail)
    '''
    {'code': 10000, 'data': 
        {'task_id': '10582841162553725037'}, 
        'message': 'Success', 
        'request_id': '20250922110356B227DB27C5C91B773F8F', 
        'status': 10000, 'time_elapsed': '247.349017ms'}
    '''
